import os
import cv2
import numpy as np
from sklearn import svm
from sklearn.metrics import precision_score, recall_score, f1_score, roc_curve, auc
import time
from matplotlib import pyplot as plt
from sklearn.utils import shuffle
from imblearn.over_sampling import SMOTE

# 定义定位胸部的函数
def locate_chest(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    _, thresh = cv2.threshold(blurred, 45, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    max_contour = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(max_contour)
    cropped_image = image[y:y+h, x:x+w]
    return cropped_image

# 加载图像并应用胸部定位函数
def load_images_from_folder(folder):
    images = []
    labels = []
    for filename in os.listdir(folder):
        for img in os.listdir(os.path.join(folder, filename)):
            img_path = os.path.join(folder, filename, img)
            image = cv2.imread(img_path)
            if image is not None:
                image = locate_chest(image)
                image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                image = cv2.resize(image, (64, 64))
                images.append(image.flatten())
                labels.append(1 if filename == 'cancer' else 0)
    return images, labels

def train_model(X_train, y_train, X_valid, y_valid):
    clf = svm.SVC(probability=True)
    clf.fit(X_train, y_train)
    y_valid_pred = clf.predict(X_valid)
    print("Validation Metrics:")
    print_metrics(y_valid, y_valid_pred)
    return clf

def predict(clf, X_test):
    y_pred = clf.predict(X_test)
    return y_pred

def print_metrics(y_test, y_pred):
    precision = precision_score(y_test, y_pred.round())
    recall = recall_score(y_test, y_pred.round())
    f1 = f1_score(y_test, y_pred.round())
    fpr, tpr, _ = roc_curve(y_test, y_pred)
    roc_auc = auc(fpr, tpr)
    print(f'Precision: {precision}\nRecall: {recall}\nF1 Score: {f1}\nAUC: {roc_auc}')
    plt.figure()
    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

def main():
    start_time = time.time()
    X_train, y_train = load_images_from_folder('D:\\kaggle\\final\\data1\\train')
    X_valid, y_valid = load_images_from_folder('D:\\kaggle\\final\\data1\\valid')
    X_test, y_test = load_images_from_folder('D:\\kaggle\\final\\data1\\test')

    # Use SMOTE to balance the data
    sm = SMOTE(random_state=42)
    X_train, y_train = sm.fit_resample(X_train, y_train)

    clf = train_model(X_train, y_train, X_valid, y_valid)
    y_pred = predict(clf, X_test)
    print("Test Metrics:")
    print_metrics(y_test, y_pred)
    end_time = time.time()
    print(f'Program running time: {end_time - start_time}')

if __name__ == "__main__":
    main()

'''
Validation Metrics:
Precision: 0.8947368421052632
Recall: 0.9577464788732394
F1 Score: 0.9251700680272109
AUC: 0.8455399061032864
Test Metrics:
Precision: 0.8701298701298701
Recall: 0.9436619718309859
F1 Score: 0.9054054054054054
AUC: 0.8051643192488265
Program running time: 180.3864815235138
'''